package io.gitee.horizon007.webspider.provider;

import com.gargoylesoftware.htmlunit.BrowserVersion;
import com.gargoylesoftware.htmlunit.WebClient;
import com.gargoylesoftware.htmlunit.html.HtmlPage;
import lombok.Setter;
import lombok.experimental.Accessors;
import org.apache.commons.collections4.MapUtils;
import org.jsoup.Jsoup;
import org.jsoup.nodes.Document;

import java.io.IOException;
import java.net.MalformedURLException;
import java.net.URL;
import java.util.Map;
import java.util.Objects;
import java.util.Optional;
import java.util.concurrent.TimeUnit;
import java.util.function.BiConsumer;
import java.util.function.Consumer;

@Setter
@Accessors(chain = true)
public class HtmlUnitProvider implements DocumentProvider, AutoCloseable {

    private int timeout = (int) TimeUnit.SECONDS.toMillis(30L);
    private BrowserVersion browserVersion;
    private Map<String, String> headers;
    private Consumer<WebClient> customizer;
    private BiConsumer<WebClient, HtmlPage> pageCustomizer;
    private long delay;

    private URL url;
    private WebClient webClient;


    @Override
    public Document getDocument() throws IOException {
        Objects.requireNonNull(url, "URL is missing");

        autoCreateWebClient();

        HtmlPage page = webClient.getPage(url);

        if (delay > 0) {
            try {
                Thread.sleep(delay);
            } catch (InterruptedException e) {
                throw new RuntimeException(e);
            }
        }

        if (pageCustomizer != null) {
            pageCustomizer.accept(webClient, page);
        }

        String content = page.asXml();

        page.cleanUp();
        page.remove();

        return Jsoup.parse(content);
    }

    @Override
    public void close() throws Exception {
        autoCloseWebClient();
    }

    private void autoCreateWebClient() {
        if (webClient == null) {
            webClient = new WebClient(Optional.ofNullable(browserVersion).orElse(BrowserVersion.getDefault()));

            webClient.getOptions().setTimeout(timeout);
            webClient.getOptions().setCssEnabled(false);
            webClient.getOptions().setDownloadImages(false);
            webClient.getOptions().setThrowExceptionOnScriptError(false);
            webClient.getOptions().setUseInsecureSSL(true);
            webClient.getOptions().setFetchPolyfillEnabled(true);
            webClient.getOptions().setProxyPolyfillEnabled(true);

            if (MapUtils.isNotEmpty(headers)) {
                for (Map.Entry<String, String> entry : headers.entrySet()) {
                    webClient.addRequestHeader(entry.getKey(), entry.getValue());
                }
            }
            if (customizer != null) {
                customizer.accept(webClient);
            }
        }
    }
    private void autoCloseWebClient() {
        if (webClient != null) {
            webClient.close();
            webClient = null;
        }
    }
}
